"""
WebSocket服务
负责实时通信，包括任务进度更新、状态通知等
"""

import logging
from typing import Dict, Any, Optional, List
from datetime import datetime
from flask_socketio import emit, join_room, leave_room, disconnect
from ..utils.redis_utils import RedisUtils

logger = logging.getLogger(__name__)

class WebSocketService:
    """WebSocket服务类"""
    
    def __init__(self, app=None, socketio=None):
        self.app = app
        self.socketio = socketio
        self.redis_utils = RedisUtils(app) if app else None
        self.connected_users = {}  # {session_id: user_info}
        
    def handle_connect(self, auth: Optional[Dict] = None) -> bool:
        """
        处理客户端连接
        
        Args:
            auth: 认证信息
            
        Returns:
            是否允许连接
        """
        try:
            from flask import request, session
            
            session_id = request.sid
            user_id = auth.get('user_id') if auth else session.get('user_id')
            
            if not user_id:
                logger.warning(f"WebSocket连接被拒绝: 缺少用户ID - {session_id}")
                return False
            
            # 记录连接信息
            self.connected_users[session_id] = {
                'user_id': user_id,
                'connected_at': datetime.utcnow(),
                'last_activity': datetime.utcnow()
            }
            
            # 加入用户房间
            join_room(f"user_{user_id}")
            
            # 发送连接成功消息
            emit('connected', {
                'status': 'success',
                'message': '连接成功',
                'timestamp': datetime.utcnow().isoformat()
            })
            
            logger.info(f"WebSocket连接成功: 用户 {user_id} - {session_id}")
            
            return True
            
        except Exception as e:
            logger.error(f"WebSocket连接处理失败: {str(e)}")
            return False
    
    def handle_disconnect(self) -> None:
        """处理客户端断开连接"""
        try:
            from flask import request
            
            session_id = request.sid
            user_info = self.connected_users.pop(session_id, None)
            
            if user_info:
                user_id = user_info['user_id']
                leave_room(f"user_{user_id}")
                
                logger.info(f"WebSocket断开连接: 用户 {user_id} - {session_id}")
            else:
                logger.info(f"WebSocket断开连接: 未知会话 - {session_id}")
                
        except Exception as e:
            logger.error(f"WebSocket断开连接处理失败: {str(e)}")
    
    def handle_join_task_room(self, data: Dict[str, Any]) -> None:
        """
        处理加入任务房间
        
        Args:
            data: 包含task_id的数据
        """
        try:
            from flask import request
            
            session_id = request.sid
            task_id = data.get('task_id')
            
            if not task_id:
                emit('error', {'message': '缺少任务ID'})
                return
            
            user_info = self.connected_users.get(session_id)
            if not user_info:
                emit('error', {'message': '用户未认证'})
                return
            
            # 验证任务权限
            if not self._verify_task_permission(task_id, user_info['user_id']):
                emit('error', {'message': '无权限访问此任务'})
                return
            
            # 加入任务房间
            join_room(f"task_{task_id}")
            
            # 发送当前任务状态
            task_status = self._get_task_status(task_id)
            if task_status:
                emit('task_status', task_status)
            
            emit('joined_task', {
                'task_id': task_id,
                'message': '已加入任务房间'
            })
            
            logger.debug(f"用户 {user_info['user_id']} 加入任务房间: {task_id}")
            
        except Exception as e:
            logger.error(f"加入任务房间失败: {str(e)}")
            emit('error', {'message': '加入任务房间失败'})
    
    def handle_leave_task_room(self, data: Dict[str, Any]) -> None:
        """
        处理离开任务房间
        
        Args:
            data: 包含task_id的数据
        """
        try:
            from flask import request
            
            session_id = request.sid
            task_id = data.get('task_id')
            
            if not task_id:
                emit('error', {'message': '缺少任务ID'})
                return
            
            user_info = self.connected_users.get(session_id)
            if not user_info:
                return
            
            # 离开任务房间
            leave_room(f"task_{task_id}")
            
            emit('left_task', {
                'task_id': task_id,
                'message': '已离开任务房间'
            })
            
            logger.debug(f"用户 {user_info['user_id']} 离开任务房间: {task_id}")
            
        except Exception as e:
            logger.error(f"离开任务房间失败: {str(e)}")
    
    def send_task_update(self, task_id: str, update_data: Dict[str, Any]) -> bool:
        """
        发送任务更新
        
        Args:
            task_id: 任务ID
            update_data: 更新数据
            
        Returns:
            是否发送成功
        """
        try:
            if not self.socketio:
                return False
            
            # 添加时间戳
            update_data['timestamp'] = datetime.utcnow().isoformat()
            update_data['task_id'] = task_id
            
            # 发送到任务房间
            self.socketio.emit('task_update', update_data, room=f"task_{task_id}")
            
            # 缓存最新状态
            if self.redis_utils:
                self.redis_utils.cache_task_status(task_id, update_data)
            
            logger.debug(f"任务更新发送: {task_id} - {update_data.get('status', 'unknown')}")
            
            return True
            
        except Exception as e:
            logger.error(f"发送任务更新失败 {task_id}: {str(e)}")
            return False
    
    def send_user_notification(self, user_id: str, notification: Dict[str, Any]) -> bool:
        """
        发送用户通知
        
        Args:
            user_id: 用户ID
            notification: 通知内容
            
        Returns:
            是否发送成功
        """
        try:
            if not self.socketio:
                return False
            
            # 添加时间戳
            notification['timestamp'] = datetime.utcnow().isoformat()
            
            # 发送到用户房间
            self.socketio.emit('notification', notification, room=f"user_{user_id}")
            
            logger.debug(f"用户通知发送: {user_id} - {notification.get('type', 'unknown')}")
            
            return True
            
        except Exception as e:
            logger.error(f"发送用户通知失败 {user_id}: {str(e)}")
            return False
    
    def broadcast_system_message(self, message: Dict[str, Any]) -> bool:
        """
        广播系统消息
        
        Args:
            message: 系统消息
            
        Returns:
            是否广播成功
        """
        try:
            if not self.socketio:
                return False
            
            # 添加时间戳
            message['timestamp'] = datetime.utcnow().isoformat()
            
            # 广播给所有连接的客户端
            self.socketio.emit('system_message', message)
            
            logger.info(f"系统消息广播: {message.get('type', 'unknown')}")
            
            return True
            
        except Exception as e:
            logger.error(f"广播系统消息失败: {str(e)}")
            return False
    
    def get_connected_users_count(self) -> int:
        """获取当前连接用户数"""
        return len(self.connected_users)
    
    def get_user_connections(self, user_id: str) -> List[str]:
        """
        获取用户的连接会话
        
        Args:
            user_id: 用户ID
            
        Returns:
            会话ID列表
        """
        return [
            session_id for session_id, user_info in self.connected_users.items()
            if user_info['user_id'] == user_id
        ]
    
    def disconnect_user(self, user_id: str, reason: str = 'server_disconnect') -> int:
        """
        断开用户的所有连接
        
        Args:
            user_id: 用户ID
            reason: 断开原因
            
        Returns:
            断开的连接数
        """
        try:
            if not self.socketio:
                return 0
            
            sessions = self.get_user_connections(user_id)
            
            for session_id in sessions:
                self.socketio.emit('disconnect_reason', {
                    'reason': reason,
                    'message': '服务器主动断开连接'
                }, room=session_id)
                
                disconnect(session_id)
                self.connected_users.pop(session_id, None)
            
            logger.info(f"断开用户连接: {user_id} - {len(sessions)} 个会话")
            
            return len(sessions)
            
        except Exception as e:
            logger.error(f"断开用户连接失败 {user_id}: {str(e)}")
            return 0
    
    def _verify_task_permission(self, task_id: str, user_id: str) -> bool:
        """验证任务权限"""
        try:
            from ..models.task import Task
            
            task = Task.find_by_id(task_id)
            return task and task.get('user_id') == user_id
            
        except Exception as e:
            logger.error(f"验证任务权限失败 {task_id}: {str(e)}")
            return False
    
    def _get_task_status(self, task_id: str) -> Optional[Dict[str, Any]]:
        """获取任务状态"""
        try:
            # 先尝试从缓存获取
            if self.redis_utils:
                cached_status = self.redis_utils.get_task_status(task_id)
                if cached_status:
                    return cached_status
            
            # 从数据库获取
            from ..models.task import Task
            
            task = Task.find_by_id(task_id)
            if task:
                return {
                    'task_id': task_id,
                    'status': task.get('status'),
                    'progress': task.get('progress', 0),
                    'message': task.get('message', ''),
                    'created_at': task.get('created_at', '').isoformat() if task.get('created_at') else None,
                    'updated_at': task.get('updated_at', '').isoformat() if task.get('updated_at') else None
                }
            
            return None
            
        except Exception as e:
            logger.error(f"获取任务状态失败 {task_id}: {str(e)}")
            return None
    
    def cleanup_inactive_connections(self, timeout_minutes: int = 30) -> int:
        """
        清理不活跃的连接
        
        Args:
            timeout_minutes: 超时分钟数
            
        Returns:
            清理的连接数
        """
        try:
            if not self.socketio:
                return 0
            
            from datetime import timedelta
            
            cutoff_time = datetime.utcnow() - timedelta(minutes=timeout_minutes)
            inactive_sessions = []
            
            for session_id, user_info in self.connected_users.items():
                if user_info['last_activity'] < cutoff_time:
                    inactive_sessions.append(session_id)
            
            # 断开不活跃的连接
            for session_id in inactive_sessions:
                user_info = self.connected_users.pop(session_id, None)
                if user_info:
                    self.socketio.emit('disconnect_reason', {
                        'reason': 'timeout',
                        'message': '连接超时'
                    }, room=session_id)
                    
                    disconnect(session_id)
            
            if inactive_sessions:
                logger.info(f"清理不活跃连接: {len(inactive_sessions)} 个")
            
            return len(inactive_sessions)
            
        except Exception as e:
            logger.error(f"清理不活跃连接失败: {str(e)}")
            return 0
    
    def update_user_activity(self, session_id: str) -> None:
        """
        更新用户活跃时间
        
        Args:
            session_id: 会话ID
        """
        user_info = self.connected_users.get(session_id)
        if user_info:
            user_info['last_activity'] = datetime.utcnow()
    
    def get_connection_statistics(self) -> Dict[str, Any]:
        """获取连接统计信息"""
        try:
            total_connections = len(self.connected_users)
            unique_users = len(set(
                user_info['user_id'] for user_info in self.connected_users.values()
            ))
            
            # 按用户统计连接数
            user_connections = {}
            for user_info in self.connected_users.values():
                user_id = user_info['user_id']
                user_connections[user_id] = user_connections.get(user_id, 0) + 1
            
            return {
                'total_connections': total_connections,
                'unique_users': unique_users,
                'user_connections': user_connections,
                'average_connections_per_user': round(total_connections / unique_users, 2) if unique_users > 0 else 0
            }
            
        except Exception as e:
            logger.error(f"获取连接统计失败: {str(e)}")
            return {
                'total_connections': 0,
                'unique_users': 0,
                'user_connections': {},
                'average_connections_per_user': 0
            }